import numpy as np
import xarray as xr
from scipy.stats import pearsonr

import matplotlib.pyplot as plt
from ppgnss import gnss_time
from ppgnss import gnss_utils

def main():
    filename = "../baseline/solution.dat"
    seal_down_file = "../data/sea_levels/downstream_7308.txt"
    seal_up_file = "../data/sea_levels/upstream_3020.txt"

    xr_seal_up = read_seal_file(seal_up_file)
    xr_seal_down = read_seal_file(seal_down_file)
    xr_seal_down, xr_seal_up = xr.align(xr_seal_down, xr_seal_up, join="inner")
    xr_delta = xr_seal_up - xr_seal_down
    fig, axes = plt.subplots(nrows=2, ncols=1,  figsize=(gnss_utils.cm2inch(10), gnss_utils.cm2inch(5)))
    xr_data = xr_seal_down.loc["2022-10-20T08:00:00":"2022-10-22T08:00:00"]
    axes[0].plot(xr_data.coords["time"].values, xr_data.values, label="down")
    ##xmin, xmax = xr_delta.loc["2022-01-01":"2023-03-20"].coords["time"][0], xr_delta.loc["2023-01-01":"2023-03-20"].coords["time"][-1]
    # axes[0].set_xlim((xmin, xmax))
    # axes[1].plot(xr_seal_up.coords["time"].values[:1000], xr_seal_up.values[:1000], label="up")
    plt.show()

    colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    year_from, doy_from = 2022, 118
    
    xr_neu = read_gamit_solution(filename, year_from, doy_from)
    # xr_neu, xr_seal_up = xr.align(xr_neu, xr_seal_up, join="inner")
    # xr_neu, xr_seal_down = xr.align(xr_neu, xr_seal_down, join="inner")
    # xr_neu, xr_seal_up = xr.align(xr_neu, xr_seal_up, join="inner")

    
    # time_arr = np.array(time_list)
    fig, ax1 = plt.subplots()
    ax1.plot(xr_neu.coords["time"], xr_neu.loc[:, "u"], label="Up", color=colors[0])
    ax1.plot(xr_neu.coords["time"], xr_neu.loc[:, "e"], label="East", color=colors[1])
    ax1.plot(xr_neu.coords["time"], xr_neu.loc[:, "n"], label="North", color=colors[2])
    ax1.set_ylim((-0.02, 0.02))

    ax2 = ax1.twinx()
    # ax2.plot(xr_seal_up.coords["time"], xr_seal_up.values, "-.", label="upup", color=colors[3])
    # ax2.plot(xr_seal_down.coords["time"], xr_seal_down.values, "-.", label="dodo", color=colors[4])
    ax2.set_ylim((-10, 10))
    lines, labels = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    plt.legend(lines + lines2, labels + labels2, loc='upper right')

    # plt.legend()
    plt.show()

    fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(gnss_utils.cm2inch(10), gnss_utils.cm2inch(5)))
    axes.plot(xr_neu.loc[:, "u"].values, xr_seal_down.values, ".", color=colors[0], label="u")
    axes.plot(xr_neu.loc[:, "e"].values, xr_seal_down.values, ".", color=colors[1], label="e")
    axes.plot(xr_neu.loc[:, "n"].values, xr_seal_down.values, ".", color=colors[2], label="n")
    axes.set_xlim((-0.005, 0.005))
    for ax in ["u", "e", "n"]:
        valid_idx = np.logical_not(np.logical_or(np.isnan(xr_neu.loc[:, ax].values), np.isnan(xr_seal_down.values))) 
        x = xr_neu.loc[:, ax].values[valid_idx]
        y = xr_seal_down.values[valid_idx]
        correlation_coefficient_u, p_value_u = pearsonr(x, y)
    
        # print(f"{ax}: ", correlation_coefficient_u, p_value_u)
        # axes[1].plot(xr_neu.loc[:, "u"].values, xr_seal_up.values, ".", color=colors[3])
        # axes[1].plot(xr_neu.loc[:, "e"].values, xr_seal_up.values, ".", color=colors[4])
        # axes[1].plot(xr_neu.loc[:, "n"].values, xr_seal_up.values, ".", color=colors[5])
    plt.legend()
    axes.set_ylim((-2, 4))
    plt.show()



def read_gamit_solution(filename, year_from, doy_from):
    _, mon_from, day_from = gnss_time.doy2ymd(year_from, doy_from)
    start_time = gnss_time.ymd2datetime(year_from, mon_from, int(day_from))
    time_coord = [_*np.timedelta64(1, "D")+np.timedelta64(12, "h")+np.datetime64(start_time) for _ in range(507)]
    neus = np.full((len(time_coord), 3), np.nan)
    xr_neu = xr.DataArray(neus, dims=["time", "data"], coords=[time_coord, ["n", "e", "u"]])
    with open(filename, "r") as fread:
        lines = fread.readlines()
        time_list = list()
        norths = list()
        easts=list()
        ups=list()
        for line in lines:
            myline = line.rstrip()
            fields = myline.split()
            # print(fields)
            year = int(fields[1][:4])
            doy = int(fields[1][-4:-1])
            north = float(fields[3])
            east = float(fields[7])
            up = float(fields[11])
            _, mon, day = gnss_time.doy2ymd(year, doy)
            strtime = f"{year:04d}-{mon:02d}-{int(day):02d} 12:00:00 "
            xr_neu.loc[strtime, "n"] = north
            xr_neu.loc[strtime, "e"] = east
            xr_neu.loc[strtime, "u"] = up


    up_ref = xr_neu[0].loc["u"].values
    east_ref = xr_neu[0].loc["e"].values
    north_ref = xr_neu[0].loc["n"].values

    dups = np.array(xr_neu.loc[:, "u"].values - up_ref)
    deasts = np.array(xr_neu.loc[:, "e"].values - east_ref)
    dnorths = np.array(xr_neu.loc[:, "n"].values - north_ref)
    invalid_idx = np.logical_or(np.logical_or(np.abs(dups) > 0.05, np.abs(dnorths) > 0.015), np.abs(deasts) > 0.02)
    invalid_idx = np.logical_or(invalid_idx, dups < -0.025)
    valid_idx = np.logical_not(invalid_idx)

    xr_neu.loc[:, "u"] = dups*valid_idx
    xr_neu.loc[:, "e"] = deasts*valid_idx
    xr_neu.loc[:, "n"] = dnorths*valid_idx
    return xr_neu


def read_seal_file(filename):
    
    with open(filename, "r") as fread:
        lines = fread.readlines()
        time_list = list()
        seals = list()
        for line in lines:
            myline = line.rstrip().strip()
            strtime = myline[:20]
            seal = float(myline[20:]) 
            if seal < -2.5:
                # pass
                seal = np.nan
            dt = gnss_time.strtime2datetime(strtime)
            dt1 = np.datetime64(dt) - np.timedelta64(8, "h")
            time_list.append(dt1)
            seals.append(seal)
        xr_seal = xr.DataArray(seals, dims=["time"], coords=[time_list])
    return xr_seal


if __name__ == "__main__":
    main()